FlagTree 前端支持

通过使用tle.barrier，tle.wgmma，tle.wgmma_wait等新增API，并基于已有tle.gpu.warp_specialize API在FlagTree Triton中实现Hopper架构上的自定义复杂流水排布。具体样例参考commit tle_hopper_fa_ws_pipelined_pingpong_persistent.py实现hopper版sdpa样例。

API修改总览：

新增 barrier API：

• tle.gpu.alloc_barriers(...)
• tle.gpu.alloc_barrier(...)
• tle.gpu.barrier_wait(bar, phaseIdx=None)
• tle.gpu.barrier_arrive(bar, arrive_count=1, phaseIdx=None)
• tle.gpu.PENDING
• tle.gpu.READY
• tle.gpu.barrier
• tle.gpu.barrier_type

新增 WGMMA API：

• tle.gpu.wgmma(...)
• tle.gpu.wgmma_wait(pendings, acc=None)

已有 API 扩展：

• tle.gpu.copy(...) 新增 barrier 参数，用于 TMA global-to-shared load 的显式 completion barrier。

Warp Specialization使用说明

使用 tle.gpu.warp_specialize API 时，需要把 producer / consumer partition 拆成独立的 @triton.jit 函数，并在 kernel 内用显式列表声明调用关系。

API 形式

tle.gpu.warp_specialize(
    functions_and_args,
    worker_num_warps,
    worker_num_regs,
)

参数：

• functions_and_args：partition 列表。每一项必须是 (fn, args)，其中 fn 是 @triton.jit 函数，args 必须是 tuple。
• functions_and_args[0]：default partition，通常作为 producer。它使用 kernel launch 的 num_warps。
• functions_and_args[1:]：worker partitions，通常作为 consumer / MMA partitions。
• worker_num_warps：每个 worker partition 的 warp 数，长度必须等于 len(functions_and_args) - 1。
• worker_num_regs：每个 worker partition 的 requested registers，长度必须等于 len(functions_and_args) - 1。

使用约定：

• functions_and_args[0] 是 default partition，通常放 producer body。

• functions_and_args[1:] 是 worker partitions，通常放 consumer / MMA body；如果需要多个 consumer，就重复添加多个 worker entry。

• 多个 consumer 需要区分行为时，可以显式传入 constexpr 参数，例如 CID。

• worker function 的参数可以捕获 shared-memory buffer、barrier、descriptor、标量 constexpr 等。多个 worker 捕获同一个值时，前端会去重并作为 ttg.warp_specialize 的 partition operand 传入。

• 使用以下 API 会启用 TLE user-promise pipeline mode：

  • tle.gpu.alloc_barriers
  • tle.gpu.alloc_barrier
  • tle.gpu.wgmma_wait

  该 mode 下：

  • 使用 mlir::triton::gpu::detail::scheduleTleWgmmaUserPromisePipeline 处理 WGMMA user-promise pipeline，并为 WGMMA 补齐 ttng::WarpGroupDotCommitOp。
  • 如果这个标记在 tle.gpu.warp_specialize 调用前已经产生，则当前 warp_specialize 会 inline 展开 functions_and_args；否则保持 tt.call 调用形式。

如果需要在 worker partition 中显式使用 barrier，应在 tle.gpu.warp_specialize 前分配 barrier，并把它作为 functions_and_args 的参数传给 producer / consumer。

  empty = tle.gpu.alloc_barrier(arrive_count=2, init=tle.gpu.READY)
  tle.gpu.warp_specialize(
      [
          (producer, (...,)),
          (consumer, (empty, 0)),
          (consumer, (empty, 1)),
      ],
      [4, 4],
      [168, 168],
  )

该样例表示重复添加两个 consumer entry，让同一段 consumer body 生成两个 worker partition。每个 worker 都使用 num_warps=4 和 requested registers=168。

用户可以在函数参数中传入 constexpr 值如 0 和 1 作为 CID 区分两个 consumer worker。上面的 empty 是 mbarrier 示例；如果要做 consumer warpgroup 之间的 ping-pong 同步，参考后面的 named barrier 示例。

Barrier 相关 API

API 参数简介

本节只说明参数含义。arrive_count 和 phaseIdx 的具体配置方式见后面的使用示例。

tle.gpu.alloc_barriers(...)

tle.gpu.alloc_barriers(
    num_barriers,
    arrive_count=1,
    init=tle.gpu.PENDING,
    expect_bytes=None,
) -> tle.gpu.barrier

参数：

• num_barriers：tle.gpu.alloc_barriers返回一个barrier数组，此处为分配的 barrier slot 数量，必须是编译期正整数。
• arrive_count：barrier 的计数参数，即 pending arrival count。必须是编译期正整数。mbarrier 和 named barrier 后端对它的解释不同。具体参考使用示例。
• init：barrier初始状态，支持 tle.gpu.PENDING 和 tle.gpu.READY。
• expect_bytes：TMA completion barrier 期望完成的字节数。为 None 或编译期正整数。

返回值：tle.gpu.barrier，表示 barrier 数组或 barrier slot。数组通过 bars[i] 取 slot。

tle.gpu.alloc_barrier(...)

tle.gpu.alloc_barrier(
    arrive_count=1,
    init=tle.gpu.PENDING,
    expect_bytes=None,
) -> tle.gpu.barrier

返回值：单个 tle.gpu.barrier slot。它是 alloc_barriers(num_barriers=1, ...) 的语法糖。

tle.gpu.barrier_wait(...)

tle.gpu.barrier_wait(bar, phaseIdx=None) -> None

参数：

• bar：要等待的 barrier slot，类型为 tle.gpu.barrier。
• phaseIdx：mbarrier logical use id。传入时使用 mbarrier 路径；为 None 时使用 named barrier 路径。

tle.gpu.barrier_arrive(...)

tle.gpu.barrier_arrive(bar, arrive_count=1, phaseIdx=None) -> None

参数：

• bar：要 arrive 的 barrier slot，类型为 tle.gpu.barrier。

• arrive_count：本次 arrive 的计数权重。

  • 只用于 mbarrier 路径，表示逻辑arrive数量；
  • named barrier 路径下必须保持默认值 1。named barrier 的总参与线程数已经由 allocation 的 arrive_count 表达。

• phaseIdx：mbarrier logical use id，第几次使用这个 slot。传入时使用 mbarrier 路径；为 None 时使用 named barrier 路径。

tle.gpu.PENDING

tle.gpu.PENDING

含义：barrier 初始不可通过。常用于 full barrier、TMA completion barrier、named barrier。

tle.gpu.READY

tle.gpu.READY

含义：barrier 初始可通过。常用于 empty barrier。READY barrier 只能走 mbarrier 路径。

tle.gpu.barrier

含义：alloc_barrier(s) 返回的 barrier value 类型。支持 bars[i] 取数组 slot。

tle.gpu.barrier_type

含义：tle.gpu.barrier 对应的类型描述。普通 kernel 代码通常不需要直接构造。

修改tle.gpu.copy(..., barrier=...)

tle.gpu.copy(src, dst, shape, offsets=None, barrier=None)

参数：

• src：copy 源。使用 barrier 时必须是 TMA TensorDescriptor。
• dst：copy 目标。使用 barrier 时必须是 shared-memory tle.gpu.buffered_tensor。
• shape：copy tile 的形状。
• offsets：TMA descriptor 坐标偏移。
• barrier：可选 completion barrier。必须来自 tle.gpu.alloc_barrier(s)(expect_bytes=...)。

使用示例

mbarrier 使用路径

传入 phaseIdx 时，barrier_wait / barrier_arrive 会选择 mbarrier 路径。tle.gpu.copy(..., barrier=...) 的 TMA completion barrier 也必须走 mbarrier 路径，并且后续 wait 必须传 phaseIdx。

• 适用场景：

  • producer/consumer Empty-Full 同步。
  • TMA global-to-shared load completion。
  • init=tle.gpu.READY 的 empty barrier。
  • 需要用 logical use id 表达同一 slot 多次复用的 ping-pong / 多 stage pipeline。

• phaseIdx 要求：

  • phaseIdx 是同一个 barrier slot 的 logical use id，也就是这个 slot 第几次被使用。它的设计意图是让循环里的同一组 barrier slot 可以被反复复用，并通过循环计数区分第几轮复用。
  • 当前实现会在内部把 logical use id 转成硬件 mbarrier phase：
    • constexpr 整数路径：(phaseIdx & 1) ^ init_polarity
    • 动态 tensor 路径：先做 phaseIdx & 1，再根据 init 调整初始 polarity
  • 用户通常传“该 slot 的复用次数”，不需要手写 & 1。
  • 同一次 logical use 对应的 wait/arrive 两侧必须传入一致的 phaseIdx。
  • 在 ring buffer / 多 stage pipeline 中，常见写法是 slot = iter % NUM_BUFFERS，phaseIdx = iter // NUM_BUFFERS，因为同一个 slot 每隔 NUM_BUFFERS 次迭代才复用一次。

  bars = tle.gpu.alloc_barriers(
      num_barriers=NUM_BUFFERS,
      arrive_count=1,
      init=tle.gpu.READY,
  )
  
  for k_iter in range(0, NUM_ITERS):
      slot = k_iter % NUM_BUFFERS
      use_id = k_iter // NUM_BUFFERS
  
      tle.gpu.barrier_wait(bars[slot], phaseIdx=use_id)
      # ... 第 use_id 次使用 bars[slot] 对应的 shared-memory slot ...
      tle.gpu.barrier_arrive(bars[slot], phaseIdx=use_id)

• arrive_count 配置示例：

  • 一个 consumer 到达同一个 slot 一次：

    @triton.jit
    def consumer(empty):
        use_id: tl.constexpr = 0
        tle.gpu.barrier_arrive(empty, phaseIdx=use_id)
    
    
    empty = tle.gpu.alloc_barrier(arrive_count=1, init=tle.gpu.READY)
    
    tle.gpu.warp_specialize(
        [
            (producer, (...,)),
            (consumer, (empty,)),
        ],
        [4],
        [168],
    )

  • 两个 worker consumers 共同释放同一个 slot，此时同一个 slot 会收到两次 arrive，所以 arrive_count=2：

    @triton.jit
    def consumer(empty):
        use_id: tl.constexpr = 0
        # 两个 worker 都执行这里，各 arrive 1 次。
        tle.gpu.barrier_arrive(empty, phaseIdx=use_id)
    
    
    empty = tle.gpu.alloc_barrier(arrive_count=2, init=tle.gpu.READY)
    
    tle.gpu.warp_specialize(
        [
            (producer, (...,)),
            (consumer, (empty,)),
            (consumer, (empty,)),
        ],
        [4, 4],
        [168, 168],
    )

  • 两个 worker consumers 使用各自独立 slot：传入barrier slots array，通过CID区分自己需要使用的slot。

    @triton.jit
    def consumer(empty, CID: tl.constexpr):
        use_id: tl.constexpr = 0
        tle.gpu.barrier_arrive(empty[CID], phaseIdx=use_id)
    
    
    empty = tle.gpu.alloc_barriers(num_barriers=2, arrive_count=1, init=tle.gpu.READY)
    
    tle.gpu.warp_specialize(
        [
            (producer, (...,)),
            (consumer, (empty, 0)),
            (consumer, (empty, 1)),
        ],
        [4, 4],
        [168, 168],
    )

• 其他注意事项：

  • 后端选择：

    • 同一个 barrier slot 不能混用 mbarrier 和 named barrier 路径。
    • barrier_wait / barrier_arrive 只要传入 phaseIdx，当前 slot 就会记录为 mbarrier backend。

  • 初始状态：

    • init=tle.gpu.READY 只适合 mbarrier 路径。
    • 如果 READY barrier 不传 phaseIdx，会选择 named barrier 路径并报错。

  • TMA completion：

    • expect_bytes barrier 只能用于 TMA global-to-shared completion。
    • 使用 expect_bytes 的 barrier 必须通过 tle.gpu.copy(..., barrier=...) 触发 completion，并且必须用 phaseIdx 等待。

named barrier 使用路径

不传 phaseIdx 时，barrier_wait / barrier_arrive 会选择 named barrier 路径。

• 适用场景：

  • 不需要 TMA completion 的 CTA-scope rendezvous。
  • consumer 内部或多个 consumer warp group 之间的轻量同步。
  • 同步点可以用静态 barrier slot 表达，且不需要 init=READY。

• named_barrier的arrive_count是参与线程数 named_num_threads ，通常写成 参与 warp 数 * 32。当前实现会把 allocation 的 arrive_count 同时传给 named arrive 和 named wait。因此，使用named_barrier只需要alloc后直接wait+arrive。举例：

• 两个 4-warp consumer workers 通过同一个 named barrier 同步，总参与线程数是 256。这里两个 worker 复用同一段 consumer body，并通过显式传入的 CID 区分一个 worker 做 non-blocking arrive，另一个 worker 做 wait/sync；不需要在 barrier_arrive/ barrier_wait调用参数里写arrive_count。

  @triton.jit
  def consumer(sync, CID: tl.constexpr):
      if CID == 0:
          tle.gpu.barrier_arrive(sync[0])
      else:
          tle.gpu.barrier_wait(sync[0])
  
  
  sync = tle.gpu.alloc_barriers(num_barriers=1, arrive_count=256)
  
  tle.gpu.warp_specialize(
      [
          (producer, (...,)),
          (consumer, (sync, 0)),
          (consumer, (sync, 1)),
      ],
      [4, 4],
      [168, 168],
  )

• 其他注意事项：

  • named barrier 只支持 init=tle.gpu.PENDING，不支持 READY。
  • named barrier 不支持 expect_bytes，因此不能用于 tle.gpu.copy(..., barrier=...) 的 TMA completion。
  • named barrier 要求静态 barrier slot index，例如 sync[0]；不能用动态 slot index。

Empty-Full 同步示例

@triton.jit
def producer(a_desc, a_smem, a_empty, a_full):
    phase: tl.constexpr = 0
    tle.gpu.barrier_wait(a_empty, phaseIdx=phase)
    tle.gpu.copy(a_desc, a_smem, [BLOCK_M, BLOCK_K], [0, 0], barrier=a_full)


@triton.jit
def consumer(a_smem, a_empty, a_full):
    phase: tl.constexpr = 0
    tle.gpu.barrier_wait(a_full, phaseIdx=phase)
    # consume a_smem ...
    tle.gpu.barrier_arrive(a_empty, phaseIdx=phase)


@triton.jit
def kernel(...):
    a_empty = tle.gpu.alloc_barrier(init=tle.gpu.READY)
    a_full = tle.gpu.alloc_barrier(expect_bytes=A_TILE_BYTES)

    tle.gpu.warp_specialize(
        [
            (producer, (a_desc, a_smem, a_empty, a_full)),
            (consumer, (a_smem, a_empty, a_full)),
        ],
        [4],
        [168],
    )

copy(..., barrier=...) 只支持 TMA global-to-shared load。不支持 shared-to-global TMA store、普通 copy、named barrier 或没有 expect_bytes 的 barrier。

TMA completion 示例

TMA completion 只能使用 mbarrier 后端，通常写 arrive_count=1，因为 arrival 来自这一条 TMA copy 的 completion：

bar = tle.gpu.alloc_barrier(expect_bytes=BLOCK_M * BLOCK_K * 2)

tle.gpu.copy(
    a_desc,
    a_smem,
    [BLOCK_M, BLOCK_K],
    [pid_m * BLOCK_M, pid_k * BLOCK_K],
    barrier=bar,
)
tle.gpu.barrier_wait(bar, phaseIdx=0)

WGMMA 相关 API

API 参数简介

tle.gpu.wgmma(...)

tle.gpu.wgmma(
    a,
    b,
    acc=None,
    input_precision=None,
    max_num_imprecise_acc=None,
    out_dtype=tl.float32,
    trans_a: tl.constexpr = False,
    trans_b: tl.constexpr = False,
) -> tl.tensor

参数：

• a：WGMMA A operand。可以是 shared-memory tle.gpu.buffered_tensor 或 register tl.tensor。
• b：WGMMA B operand。当前必须是 shared-memory tle.gpu.buffered_tensor。
• acc：输入 accumulator。为 None 时前端会生成零 accumulator。
• input_precision：dot input precision。为 None 时使用当前 builder 默认值。
• max_num_imprecise_acc：FP8 imprecise accumulation 控制参数，不改变返回 tensor dtype，仅在FP8 GEMM上有效。使用 FP8 -> F32 WGMMA 当前有效值需 >= 32；None 在 sm90 上默认 2^30。值越小越偏精度，越大越偏性能。
• out_dtype：输出 accumulator dtype，默认 tl.float32。它决定结果/accumulator 的 dtype；
• trans_a：是否转置 A operand，必须是编译期 bool。
• trans_b：是否转置 B operand，必须是编译期 bool。

返回值：异步 WGMMA accumulator dependency value。普通 tensor op、reduce、store 消费前必须调用 tle.gpu.wgmma_wait(...)。

tle.gpu.wgmma_wait(...)

tle.gpu.wgmma_wait(pendings, acc=None) -> tl.tensor

参数：

• pendings：允许保留的 outstanding WGMMA group 数量。0 表示完全等待。
• acc：要等待的 accumulator。若省略且第一个参数是 tensor，则等价于 tle.gpu.wgmma_wait(0, acc)。

返回值：wait 后可安全消费的 accumulator tensor。

使用示例

单次 WGMMA：

a_smem = tle.gpu.alloc([64, 16], dtype=tl.float16, layout=None, scope=tle.gpu.smem)
b_smem = tle.gpu.alloc([16, 16], dtype=tl.float16, layout=None, scope=tle.gpu.smem)

acc = tle.gpu.wgmma(a_smem, b_smem, out_dtype=tl.float32)
acc = tle.gpu.wgmma_wait(0, acc)

流水中保留一个 pending group，WGMMA是warpgroup级别的FIFO队列模型：

acc1 = tle.gpu.wgmma(a0_smem, b0_smem, acc1)
acc2 = tle.gpu.wgmma(a1_smem, b1_smem, acc2)
## do something else
acc1 = tle.gpu.wgmma_wait(1, acc1) # permit acc2 is pending

当前性能

测试说明

• vllm：0.17.0
  • 测试函数：vllm.vllm_flash_attn.flash_attn_interface.flash_attn_varlen_func
• torch：2.11.0a0+eb65b36914.nv26.02（docker：http://nvcr.io/nvidia/pytorch:26.02-py3）
  • 测试函数：torch.nn.attention.varlen.varlen_attn
• GPU：H800
  • NVIDIA-SMI 570.124.06
  • Driver Version: 570.124.06
  • CUDA Version: 13.1

通过FlagGems benchmark测试，右侧三列表示以vllm为100%性能，其他三个的性能比。